package com.rt.gptnext.common.utils.http.engine;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.ssl.TrustAnyHostnameVerifier;
import cn.hutool.json.JSONUtil;
import com.rt.gptnext.common.constant.HttpConstants;
import com.rt.gptnext.common.utils.SafeUtil;
import com.rt.gptnext.common.utils.http.HttpEngine;
import com.rt.gptnext.common.utils.http.HttpInfo;
import com.rt.gptnext.common.utils.http.HttpRequest;
import com.rt.gptnext.common.utils.http.HttpStream;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;


/**
 * @author gaoyiyang
 **/
public class OkhttpEngine extends HttpEngine {

    private final OkHttpClient client;

    public OkhttpEngine(boolean pool, int maxIdleConnections, int timeoutSeconds) {
        OkHttpClient.Builder builder = new OkHttpClient.Builder();
        if (pool) {
            builder.connectionPool(new ConnectionPool(maxIdleConnections, 10, TimeUnit.MINUTES));
        }
        client = builder.callTimeout(timeoutSeconds, TimeUnit.SECONDS)
                .connectTimeout(timeoutSeconds, TimeUnit.SECONDS)
                .readTimeout(timeoutSeconds, TimeUnit.SECONDS)
                .writeTimeout(timeoutSeconds, TimeUnit.SECONDS)
                .sslSocketFactory(createSslSocketFactory(), new MyTrustManager())
                .hostnameVerifier(new TrustAnyHostnameVerifier())
                .build();
    }

    @Override
    protected void requestImpl(HttpRequest request, HttpInfo info) {
        try {
            Request req = createRequest(request, info);
            info.setRequestHeaders(req.headers().toMultimap());
            // 发送请求
            try (Response resp = client
                    .newCall(req)
                    .execute()) {
                info.setStatusCode(resp.code());
                info.setResponseHeaders(resp.headers().toMultimap());
                info.setResponseBody(Objects.requireNonNull(resp.body()).bytes());
                // 如果设置了输出流，异步将结果写入流中
                if (request.getOut() != null) {
                    ThreadUtil.execAsync(() -> IoUtil.write(request.getOut(), true, info.getResponseBody()));
                }
            }
        } catch (Throwable e) {
            info.setError(e);
            log().error("请求失败", e);
        }
    }

    @Override
    protected void requestStreamImpl(HttpRequest request, Consumer<HttpStream> consumer, CountDownLatch block, HttpInfo info, Runnable end) {
        Request req = createRequest(request, info);
        info.setRequestHeaders(req.headers().toMultimap());
        info.setDownload(request.isDownload());
        // 发送请求
        client.newCall(req).enqueue(new Callback() {
            @Override
            public void onFailure(@NotNull Call call, @NotNull IOException e) {
                try {
                    info.setError(e);
                    log().error("请求失败", e);
                    info.setSuccess(false);
                } finally {
                    end.run();
                    if (block != null) {
                        block.countDown();
                    }
                }

            }

            @Override
            public void onResponse(@NotNull Call call, @NotNull Response resp) {
                try (resp) {
                    info.setStatusCode(resp.code());
                    info.setResponseHeaders(resp.headers().toMultimap());
                    if (resp.code() >= 200 && resp.code() < 300 && resp.body() != null) {
                        var reader = resp.body().byteStream();
                        consumer.accept(new HttpStreamImpl(reader));
                        info.setResponseBody("[stream]".getBytes());
                    } else if (resp.body() != null) {
                        var string = resp.body().string();
                        consumer.accept(new HttpStreamImpl(IoUtil.toStream(string, StandardCharsets.UTF_8)));
                        info.setResponseBody(string.getBytes());
                    }
                    // 校验请求是否成功
                    info.setSuccess(info.getStatusCode() >= 200 && info.getStatusCode() < 300);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                } finally {
                    end.run();
                    if (block != null) {
                        block.countDown();
                    }
                }
            }
        });
    }

    private Request createRequest(HttpRequest request, HttpInfo info) {
        // 处理url参数
        String url = handleUrl(request, info);
        // 创建请求体
        RequestBody body = createRequestBody(request, info);
        // 构造okhttp请求对象
        Request.Builder builder = new Request
                .Builder()
                .url(url)
                .method(request.getMethod(), body);
        // 添加请求头
        if (CollUtil.isNotEmpty(request.getHeaders())) {
            request.getHeaders().forEach(builder::addHeader);
        }
        return builder.build();
    }

    /**
     * 构造okhttp请求体
     *
     * @param request 请求对象
     * @param info    请求信息
     * @return okhttp请求体
     */
    private RequestBody createRequestBody(HttpRequest request, HttpInfo info) {
        info.setRequestBody(request.getBody());
        RequestBody body = null;
        if (request.getBody() != null) {
            if (JSONUtil.isTypeJSON(request.getBody())) {
                body = RequestBody.create(request.getBody(), MediaType.parse(HttpConstants.MimeTypes.JSON_UTF8));
            } else {
                body = RequestBody.create(request.getBody().getBytes());
            }
        } else if (CollUtil.isNotEmpty(request.getMultipartList())) {
            List<String> requestBody = new ArrayList<>();
            // 多请求体表单
            MultipartBody.Builder multipartBodyBuilder = new MultipartBody.Builder("---WebKitFormBoundary7MA4YWxkTrZu0gW");
            multipartBodyBuilder.setType(Objects.requireNonNull(MediaType.parse("multipart/form-data")));
            request.getMultipartList().forEach((multipart) -> {
                if (multipart.isFile() && multipart.getFileName() == null) {
                    return;
                }
                RequestBody partBody = RequestBody.create(multipart.getValue(), MediaType.parse(multipart.getContentType()));
                if (multipart.isFile()) {
                    requestBody.add(multipart.getName() + "=[file:" + multipart.getFileName() + "]");
                    multipartBodyBuilder.addFormDataPart(multipart.getName(), multipart.getFileName(), partBody);
                } else {
                    requestBody.add(multipart.getName() + "=" + new String(multipart.getValue()).replace("\r", "\\r").replace("\n", "\\n"));
                    multipartBodyBuilder.addFormDataPart(multipart.getName(), null, partBody);
                }
            });
            body = multipartBodyBuilder.build();
            info.setRequestBody(CollUtil.join(requestBody, ";"));
        }
        if (body == null && !StrUtil.equals(request.getMethod(), HttpConstants.GET)) {
            body = RequestBody.create(new byte[0]);
        }
        return body;
    }

    private SSLSocketFactory createSslSocketFactory() {
        SSLSocketFactory ssfFactory = null;
        try {
            MyTrustManager mMyTrustManager = new MyTrustManager();
            SSLContext sc = SSLContext.getInstance("TLS");
            sc.init(null, new TrustManager[]{mMyTrustManager}, new SecureRandom());
            ssfFactory = sc.getSocketFactory();
        } catch (Exception ignored) {
        }
        return ssfFactory;
    }

    /**
     * 实现X509TrustManager接口
     */
    public static class MyTrustManager implements X509TrustManager {
        @Override
        public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {
        }

        @Override
        public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException {
        }

        @Override
        public X509Certificate[] getAcceptedIssuers() {
            return new X509Certificate[0];
        }
    }

    public static class HttpStreamImpl implements HttpStream {

        private final InputStream in;

        private final BufferedReader reader;

        private HttpStreamImpl(InputStream in) {
            this.in = in;
            this.reader = IoUtil.getReader(in, StandardCharsets.UTF_8);
        }

        @Override
        public String readLine() {
            if (this.reader == null) {
                return null;
            }
            return SafeUtil.call(reader::readLine);
        }

        @Override
        public String readStr() {
            if (this.in == null) {
                return null;
            }
            try {
                byte[] buf = new byte[1024];
                var len = in.read(buf);
                if (len <= 0) {
                    return null;
                }
                return StrUtil.str(Arrays.copyOf(buf, len), StandardCharsets.UTF_8);
            } catch (Exception ignored) {
            }
            return null;
        }
    }

}
